package utils

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"time"

	"github.com/golang-jwt/jwt/v4"
)

type JWTManager struct {
	secretKey            string
	pubKey               string
	tokenDuration        time.Duration
	refreshTokenDuration time.Duration
}

type UserClaims struct {
	UserID   int64  `json:"userId"`
	Name     string `json:"name"`
	Status   int8   `json:"status"`
	UserType int64  `json:"userType"`
	jwt.RegisteredClaims
}

func NewJWTManager(param NewJWTManagerParam) *JWTManager {
	if param.TokenDuration == 0 {
		param.TokenDuration = time.Duration(TokenExpired)
	}

	if param.RefreshTokenDuration == 0 {
		param.RefreshTokenDuration = time.Duration(RefreshTokenExpired)
	}

	if param.SecretKey == "" || param.PubKey == "" {
		param.SecretKey = privateKey
		param.PubKey = publicKey
	}

	return &JWTManager{
		secretKey:            param.SecretKey,
		pubKey:               param.PubKey,
		tokenDuration:        param.TokenDuration,
		refreshTokenDuration: param.RefreshTokenDuration,
	}
}

func (m *JWTManager) Generate(param GenerateParam) (string, string, error) {
	dePrivateKey, err := DecodeKey([]byte(privateKey))
	if err != nil {
		return "", "", ErrTokenAuthErr
	}

	claims := UserClaims{
		UserID:   param.UserID,
		Name:     param.Name,
		Status:   param.Status,
		UserType: param.UserType,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    "TUD0",
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(m.tokenDuration)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
		},
	}

	rClaims := claims
	rClaims.RegisteredClaims.ExpiresAt = jwt.NewNumericDate(time.Now().Add(m.refreshTokenDuration))
	tokerStr, err := jwt.NewWithClaims(jwt.SigningMethodRS256, claims).SignedString(dePrivateKey)
	if err != nil {
		return "", "", err
	}

	rtokerStr, err := jwt.NewWithClaims(jwt.SigningMethodRS256, rClaims).SignedString(dePrivateKey)
	if err != nil {
		return "", "", err
	}

	return tokerStr, rtokerStr, nil
}

func (m *JWTManager) Refresh(refreshToken string) (string, string, error) {
	token, err := jwt.ParseWithClaims(
		refreshToken,
		&UserClaims{},
		func(t *jwt.Token) (interface{}, error) {
			return []byte(m.secretKey), nil
		},
	)

	if err != nil {
		return "", "", err
	}

	var cliams *UserClaims
	if val, ok := token.Claims.(*UserClaims); ok && token.Valid {
		cliams = val
	} else {
		return "", "", ErrInvalidToken
	}

	// 生成新的token对
	accessToken, refreshToken, err := m.Generate(GenerateParam{
		UserID:   cliams.UserID,
		Name:     cliams.Name,
		UserType: cliams.UserType,
		Status:   cliams.Status,
	})
	if err != nil {
		return "", "", ErrTokenAuthErr
	}

	return accessToken, refreshToken, nil
}

func (m *JWTManager) Verify(tokenStr string) (*UserClaims, error) {
	pubKey, err := DecodePubKey([]byte(m.pubKey))
	if err != nil {
		return nil, err
	}

	token, err := jwt.ParseWithClaims(
		tokenStr,
		&UserClaims{},
		func(t *jwt.Token) (interface{}, error) {
			return pubKey, nil
		},
	)

	if err != nil {
		return nil, err
	}

	if claims, ok := token.Claims.(*UserClaims); ok && token.Valid {
		return claims, nil
	}

	return nil, jwt.ErrInvalidKey
}

func DecodeKey(data []byte) (*rsa.PrivateKey, error) {
	signKey, err := jwt.ParseRSAPrivateKeyFromPEM(data)
	if err != nil {
		return nil, err
	}

	return signKey, nil
}

func DecodePubKey(key []byte) (*rsa.PublicKey, error) {
	p := &pem.Block{}
	p, _ = pem.Decode(key)
	if p == nil {
		return nil, errors.New("parse key error")
	}
	return x509.ParsePKCS1PublicKey(p.Bytes)
}
